import csv
import re
import sys
import logging


def process_school_data(input_file, output_file):
    try:
        # 配置日志
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

        # 读取CSV文件
        with open(input_file, 'r', encoding='utf-8-sig') as file:
            reader = csv.DictReader(file)
            data = list(reader)

        # 标准化列名
        renamed_data = []
        for row in data:
            renamed_row = {}
            for key, value in row.items():
                if key == '序号':
                    renamed_row['serial_no'] = value
                elif key in ['学校名称', '学校']:
                    renamed_row['school_name'] = value
                elif key == '学校性质':
                    renamed_row['school_type'] = value
                elif key == '教职工人数':
                    renamed_row['total_staff_count'] = value
                elif key == '专职教师人数':
                    renamed_row['full_time_teacher_count'] = value
                elif key == '师生比':
                    renamed_row['student_faculty_ratio_raw'] = value
                else:
                    renamed_row[key] = value  # 保留其他列
            renamed_data.append(renamed_row)

        # 筛选公办学校
        public_schools = [row for row in renamed_data if row.get('school_type') == '公办']

        # 清理师生比列
        valid_ratio_pattern = re.compile(r'^1:[0-9.]+$')
        valid_ratio_data = []
        for row in public_schools:
            ratio_str = row.get('student_faculty_ratio_raw', '')
            if ratio_str and valid_ratio_pattern.match(ratio_str):
                valid_ratio_data.append(row)

        # 提取比率数值部分并计算学生数
        processed_data = []
        for row in valid_ratio_data:
            try:
                ratio_str = row['student_faculty_ratio_raw']
                match = re.search(r'^1:([0-9.]+)$', ratio_str)
                if match:
                    ratio_value = float(match.group(1))

                    # 尝试转换教职工人数为数值
                    try:
                        staff_count = float(row.get('total_staff_count', 0))
                    except ValueError:
                        staff_count = 0

                    estimated_students = staff_count * ratio_value

                    # 添加计算结果
                    new_row = row.copy()
                    new_row['estimated_student_count'] = estimated_students
                    processed_data.append(new_row)
            except Exception as e:
                logging.error(f"处理行时出错: {e}, 行数据: {row}")

        # 处理缺失数据并确保数值类型
        final_data = []
        required_columns = ['serial_no', 'school_name', 'total_staff_count',
                            'full_time_teacher_count', 'estimated_student_count']

        for row in processed_data:
            valid = True
            for col in required_columns:
                if col not in row or not row[col]:
                    valid = False
                    break
            if valid:
                # 确保数值类型
                try:
                    row['total_staff_count'] = float(row['total_staff_count'])
                    row['full_time_teacher_count'] = float(row['full_time_teacher_count'])
                    row['estimated_student_count'] = float(row['estimated_student_count'])
                    final_data.append(row)
                except ValueError as e:
                    logging.error(f"转换数值类型时出错: {e}, 行数据: {row}")

        # 选择并排序最终列
        final_output = []
        for row in final_data:
            output_row = {
                'serial_no': row['serial_no'],
                'school_name': row['school_name'],
                'total_staff_count': row['total_staff_count'],
                'full_time_teacher_count': row['full_time_teacher_count'],
                'estimated_student_count': row['estimated_student_count']
            }
            final_output.append(output_row)

        # 保存输出
        with open(output_file, 'w', newline='', encoding='utf-8-sig') as file:
            writer = csv.DictWriter(file, fieldnames=required_columns)
            writer.writeheader()
            writer.writerows(final_output)

        logging.info(f"数据处理完成，已保存至 {output_file}")
        return True

    except Exception as e:
        logging.error(f"处理数据时发生错误: {e}")
        return False


if __name__ == "__main__":
    # 默认输入输出文件
    input_file = 'baoshan-schools-2025.csv'
    output_file = 'clean-baoshan-schools-2025.csv'

    # 检查是否有命令行参数
    if len(sys.argv) > 1:
        input_file = sys.argv[1]
    if len(sys.argv) > 2:
        output_file = sys.argv[2]

    # 处理数据
    process_school_data(input_file, output_file)